from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *
import sys
sys.path.append("../")
from drivers.driver_serial import *
from ui.Ui_Serial_Widget import Ui_Serial_Qwidget
import threading
from ui.MyFigure import MyFigure
import time
import csv
import os
from drivers.cimoread import *
import numpy as np
import serial
import matplotlib.pyplot as plt 

# 设置matplotlib支持中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False  # 解决负号显示问题

class SensorDataCollector:
    def __init__(self, port, baud_rate, command_hex, expected_length, packet_length, file_path):
        self.port = port
        self.baud_rate = baud_rate
        self.command_hex = command_hex
        self.expected_length = expected_length
        self.packet_length = packet_length
        self.file_path = file_path
        self.ser = None
        self.connect()
    
    def connect(self):
        try:
            self.ser = serial.Serial(self.port, self.baud_rate, timeout=1)
        except Exception as e:
            print(f"串口连接错误: {e}")
    
    def collect_sensor_data(self):
        # 模拟返回传感器数据
        return np.random.rand(7) * 100
    
    def close(self):
        if self.ser and self.ser.is_open:
            self.ser.close()


# 数据采集工作线程
class DataCollectionThread(QThread):
    data_updated = pyqtSignal(list, list, list, list, int)  # 新增批次号参数
    reset_chart = pyqtSignal(int)  # 重置图表信号，附带批次号  
      
    def __init__(self, device_name, band, parent=None):
        super().__init__(parent)
        self.device_name = device_name
        self.band = band
        self.start_event = threading.Event()
        self.stop_event = threading.Event()
        self.sensor_collector = None
        self.csv_file = None
        self.csv_writer = None
        self.data_mean = [0] * 7
    
    def run(self):
        try:
            self.sensor_collector = self.initialize_components()
            self.setup_csv()
            
            current_batch = 0  # 记录当前是第几批数据
            
            while not self.stop_event.is_set():
                self.start_event.wait()
                
                # 每批数据独立的存储
                x_data = []
                y_data_x = []
                y_data_y = []
                y_data_z = []
                
                try:
                    # 发送重置信号，开始新批次
                    self.reset_chart.emit(current_batch)
                    
                    while not self.stop_event.is_set() and len(x_data) < 100:
                        self.start_event.wait()
                        
                        data = self.sensor_collector.collect_sensor_data()
                        data = [data[i] - self.data_mean[i] for i in range(7)]
                    
                        if len(data) >= 3:
                            x_val, z_val, y_val = data[:3]
                            
                            # 保存原始数据，x坐标从0开始
                            x_data.append(len(x_data))
                            y_data_x.append(x_val)
                            y_data_y.append(y_val)
                            y_data_z.append(z_val)
                            
                            # 发送信号更新UI，带上批次号
                            self.data_updated.emit(x_data, y_data_x, y_data_y, y_data_z, current_batch)
                            
                            # 写入CSV (添加批次信息)
                            if self.csv_writer:
                                self.csv_writer.writerow([current_batch, len(x_data), x_val, y_val, z_val])
                        
                        self.msleep(100)  # 控制采样频率
                    
                    # 增加批次计数
                    current_batch += 1
                
                except Exception as e:
                    print(f"数据处理错误: {e}")
                    
        except Exception as e:
            print(f"线程运行错误: {e}")
        finally:
            self.cleanup()
    
    def initialize_components(self):
        # 初始化传感器采集器并获取平均值
        sensor_collector = SensorDataCollector(
            port=self.device_name,
            baud_rate=self.band,
            command_hex='5A000603035D',
            expected_length=348,
            packet_length=58,
            file_path="1.csv"
        )
        
        sum_data = [0.0] * 7
        for _ in range(10):
            data = sensor_collector.collect_sensor_data()
            for i in range(7):
                sum_data[i] += data[i]
        
        self.data_mean = [x / 10 for x in sum_data]
        return sensor_collector
    
    def setup_csv(self):
        csv_filename = "force_data.csv"
        if not os.path.exists(csv_filename):
            with open(csv_filename, mode='w', newline='') as file:
                writer = csv.writer(file)
                # 修改表头以匹配实际数据
                writer.writerow(["Batch", "Index", "ForceX", "ForceY", "ForceZ"])
        
        self.csv_file = open(csv_filename, mode='a', newline='')
        self.csv_writer = csv.writer(self.csv_file)
    
    def cleanup(self):
        if self.csv_file:
            self.csv_file.close()
        if self.sensor_collector:
            self.sensor_collector.close()
    
    def start_collection(self):
        self.start_event.set()
    
    def stop_collection(self):
        self.start_event.clear()
        self.stop_event.set()
        self.wait()  # 等待线程完成

# 主窗口类
class MyWidget(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.ui = Ui_Serial_Qwidget()
        self.ui.setupUi(self)   
        self.device = ()
        self.device_name = "COM4"
        self.band = 0
        self.run_DataStart = None
        self.F1 = None
        self.init_ui()
    
    def Refresh(self):
        # 模拟扫描串口
        self.device = [("COM4", "虚拟串口")]
        if self.device:
            print("可用串口:")
            for device, description in self.device:
                print(f"{device} --> {description}")
        else:
            print("未找到串口")
            
        device_name_0 = [i[0] for i in self.device]
        self.ui.comboBox.clear()
        self.ui.comboBox.addItems(device_name_0)

    def on_connect_clicked(self):
        if self.device:
            self.device_name = self.ui.comboBox.currentText()
            self.band = self.ui.comboBox_2.currentText()
            print(f"连接到 {self.device_name}, 波特率 {self.band}")
        else:
            print("无可用端口！")
    
    def plot_and_update(self, x, y_x, y_y, y_z, batch_number):
        if not self.F1:
            self.run_FigureStart_assist()
        if self.F1:
            self.F1.plotFig(x, y_x, y_y, y_z)
            # 更新批次号显示
            self.setWindowTitle(f"力传感器数据采集 - 批次 {batch_number}")
    
    def run_FigureStart_assist(self):
        self.F1 = MyFigure(width=5, height=4, dpi=100)
        width, height = self.ui.graphicsView.width(), self.ui.graphicsView.height()
        self.F1.setFixedSize(width, height)
        self.scene = QGraphicsScene()
        self.scene.addWidget(self.F1)
        self.ui.graphicsView.setScene(self.scene)
    
    def DataStart(self):
        if not self.device_name or not self.band:
            print("请先选择端口和波特率")
            return
        
        if not self.run_DataStart or not self.run_DataStart.isRunning():
            self.run_DataStart = DataCollectionThread(self.device_name, self.band)
            self.run_DataStart.data_updated.connect(self.plot_and_update)
            self.run_DataStart.reset_chart.connect(self.reset_chart)  # 连接重置信号
            self.run_DataStart.start()
            self.run_DataStart.start_collection()
            
            # 开启图像显示
            if not self.F1:
                self.run_FigureStart_assist()
    
    def DataStop(self):
        if self.run_DataStart and self.run_DataStart.isRunning():
            self.run_DataStart.stop_collection()
            self.run_DataStart = None
    
    def reset_chart(self, batch_number):
        # 清空图表并更新批次号
        if self.F1:
            self.F1.clear_chart(batch_number)
            print(f"开始新的批次: {batch_number}")
    
    def init_ui(self):
        self.ui.pushButton.clicked.connect(self.Refresh)
        self.ui.pushButton_2.clicked.connect(self.on_connect_clicked)
        self.ui.pushButton_3.clicked.connect(self.DataStart)
        self.ui.pushButton_4.clicked.connect(self.DataStop)

if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MyWidget()
    window.show()
    sys.exit(app.exec_())